import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from heapq import heappush, heappop

# 论文参数配置
a = 0.71  # 体力消耗公式参数
b = 0.51  # 体力消耗公式参数
c_up = 2.7  # 上楼体力系数
c_down = 1.0  # 下楼体力系数
base_speed = 1.0  # 基础速度(m/s)
stair_height = 3.0  # 单层楼梯高度(m)
cell_size = 0.6  # 元胞尺寸(m)

class Building:
    def __init__(self, floors=40, size=20):
        self.floors = floors  # 总层数
        self.size = size      # 每层平面尺寸
        self.stair_locations = [(size//2, size//2)] * floors  # 楼梯中心位置
        self.exit_location = (0, 0, 0)  # 最终出口位置

    def get_neighbors(self, x, y, z):
        """36邻域三维邻居生成"""
        neighbors = []
        for dz in [-1, 0, 1]:
            for dx in [-1, 0, 1]:
                for dy in [-1, 0, 1]:
                    if dx == dy == dz == 0:
                        continue
                    if abs(dx) + abs(dy) + abs(dz) <= 2:
                        neighbors.append((x+dx, y+dy, z+dz))
        return neighbors

class Person:
    def __init__(self, floor, position):
        self.x, self.y = position
        self.z = floor
        self.speed = np.random.normal(1.0, 0.1)  # 初始速度正态分布
        self.total_distance = 0.0
        self.total_time = 0.0
        self.step_counter = 0
        self.step_required = int(1 / self.speed)  # 时间步长计算

    def update_speed(self, is_ascending):
        """体力消耗速度更新公式"""
        c = c_up if is_ascending else c_down
        speed = base_speed * (1 - 0.01 * c * (self.total_distance**a) * (self.total_time**b))
        self.speed = max(speed, 0.3)  # 保持最小速度
        self.step_required = int(1 / self.speed)

class EvacuationModel:
    def __init__(self, building, population=400):
        self.building = building
        self.population = [Person(z//10, (np.random.randint(5,15), np.random.randint(5,15))) 
                         for z in range(building.floors) for _ in range(10)]
        
        # 静态场计算（Dijkstra算法）
        self.static_field = np.full((building.size, building.size, building.floors), np.inf)
        self.calculate_static_field()

    def calculate_static_field(self):
        """三维静态场计算"""
        heap = []
        exit_x, exit_y, exit_z = self.building.exit_location
        self.static_field[exit_x][exit_y][exit_z] = 0
        heappush(heap, (0, exit_x, exit_y, exit_z))

        while heap:
            dist, x, y, z = heappop(heap)
            for dx, dy, dz in self.building.get_neighbors(x, y, z):
                if 0 <= dx < self.building.size and 0 <= dy < self.building.size and 0 <= dz < self.building.floors:
                    new_dist = dist + np.sqrt(dx**2 + dy**2 + (dz*stair_height)**2)
                    if new_dist < self.static_field[dx][dy][dz]:
                        self.static_field[dx][dy][dz] = new_dist
                        heappush(heap, (new_dist, dx, dy, dz))

    def evacuation_step(self):
        """疏散模拟步进"""
        moved = []
        for person in self.population:
            # 更新体力消耗
            person.step_counter += 1
            if person.step_counter >= person.step_required:
                person.total_time += 1
                person.total_distance += cell_size
                is_ascending = person.z < self.building.exit_location[2]  # 判断上下楼
                person.update_speed(is_ascending)
                
                # 选择最优移动方向
                neighbors = self.building.get_neighbors(person.x, person.y, person.z)
                best_move = min(neighbors, key=lambda p: self.static_field[p[0]%self.building.size, 
                                                                          p[1]%self.building.size, 
                                                                          p[2]%self.building.floors])
                
                # 楼梯转移逻辑
                if (best_move[0], best_move[1]) == self.building.stair_locations[person.z]:
                    if person.z > 0:
                        best_move = (best_move[0], best_move[1], person.z-1)
                
                # 执行移动
                person.x, person.y, person.z = best_move
                person.step_counter = 0
                moved.append(person)
        
        # 可视化更新
        self.update_visualization(moved)

    def update_visualization(self, moved):
        """三维可视化更新"""
        fig = plt.figure(figsize=(10, 8))
        ax = fig.add_subplot(111, projection='3d')
        
        # 绘制建筑结构
        z = np.arange(self.building.floors)
        x, y = np.meshgrid(range(self.building.size), range(self.building.size))
        ax.plot_surface(x, y, np.zeros_like(x), alpha=0.3)  # 地面
        
        # 绘制人员
        for p in self.population:
            ax.scatter(p.x, p.y, p.z, c='red' if p in moved else 'blue', s=50)
        
        # 绘制楼梯
        for z in range(self.building.floors):
            sx, sy = self.building.stair_locations[z]
            ax.scatter(sx, sy, z, c='green', s=100, marker='s')
        
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Floor')
        plt.show()

# 初始化模拟
building = Building(floors=40, size=20)
model = EvacuationModel(building)

# 运行模拟
for _ in range(100):
    model.evacuation_step()
